#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy, sys
import moveit_commander
import math
from std_msgs.msg import Int8
from xarm_moveit_demo.srv import *
from geometry_msgs.msg import PoseStamped, Pose
from moveit_commander import MoveGroupCommander, PlanningSceneInterface
from moveit_msgs.msg import PlanningScene, ObjectColor
from moveit_msgs.msg import Grasp,PlaceLocation,GripperTranslation, MoveItErrorCodes
from trajectory_msgs.msg import JointTrajectory, JointTrajectoryPoint
from tf.transformations import quaternion_from_euler
from copy import deepcopy

GROUP_NAME_ARM = 'xarm'
GROUP_NAME_GRIPPER = 'gripper'
GRIPPER_FRAME = 'gripper_centor_link'
GRIPPER_OPEN = [0.65,0.65]
GRIPPER_GRASP = [0.1,0.1]
GRIPPER_CLOSED = [0.0,0.0]
GRIPPER_JOINT_NAMES = ['gripper_1_joint','gripper_2_joint']
GRIPPER_EFFORT = [1.0,1.0]
REFERENCE_FRAME = 'base_link'

class MoveItPickPlaceDemo:
    def __init__(self):
        # 初始化Python API 依赖的moveit_commanderC++系统
        moveit_commander.roscpp_initialize(sys.argv)
        # 初始化节点
        rospy.init_node('moveit_pick_place_demo')

        # 初始化场景对象，用来在规划场景中添加或移除物体
        self.scene = PlanningSceneInterface()
        rospy.sleep(1)

        # 创建一个话题发布端用来将抓取位姿发布到gripper_pose上，在RViz中可视化
        self.gripper_pose_pub = rospy.Publisher('gripper_pose', PoseStamped,queue_size=10)

        # 定义ROS服务/pick_place_demo的服务端，服务回调函数为call_pick_place()
        self.pick_place_srv = rospy.Service('/pick_place_demo', CallPickPlaceDemo, self.call_pick_place)

        # Initialize the move group for the xarm
        self.xarm = MoveGroupCommander(GROUP_NAME_ARM)

        # Initialize the move group for the gripper
        self.gripper = MoveGroupCommander(GROUP_NAME_GRIPPER)
        rospy.loginfo("Pick and Place demo is ready. You can call the service /pick_place_demo to test... ")

        rospy.spin()

        # Shut down MoveIt cleanly
        moveit_commander.roscpp_shutdown()

        # Exit the script
        moveit_commander.os._exit(0)


    # /pick_place_demo的服务回调函数
    def call_pick_place(self, req):
        # 获取 end-effector link的名字
        end_effector_link = self.xarm.get_end_effector_link()

        # 允许重规划
        self.xarm.allow_replanning(True)

        # 设置目标的参考系
        self.xarm.set_pose_reference_frame(REFERENCE_FRAME)

        # 设置每次规划尝试的时间为5s
        self.xarm.set_planning_time(5)

        # 定义抓取操作的最大尝试次数
        max_pick_attempts = 5

        # 定义放置操作的最大尝试次数
        max_place_attempts = 5

        # 设定规划场景里的物体的名字标记，桌子table和目标物体target
        table_id = 'table'
        target_id = 'target'

        # 分离运行上一个代码后可能附着到机械臂上的目标物体
        self.scene.remove_attached_object(GRIPPER_FRAME, target_id)
        # 移除运行上一个代码后可能在规划场景里添加的桌子和目标物体
        self.scene.remove_world_object(table_id)
        self.scene.remove_world_object(target_id)
        # 给规划场景一定的更新时间
        rospy.sleep(1)

        # 机械臂先回到初始位置
        self.xarm.set_named_target('Home')
        self.xarm.go()
        rospy.sleep(1)

        # 设置桌子的长宽高 [l, w, h]
        table_size = [1.0, 1.2, 0.01]
        # 设置桌子的位姿
        table_pose = PoseStamped()
        table_pose.header.frame_id = REFERENCE_FRAME
        table_pose.pose.position.x = 0.0
        table_pose.pose.position.y = 0.0
        table_pose.pose.position.z = - table_size[2] / 2.0
        table_pose.pose.orientation.w = 1.0
        # 把桌子添加到规划场景中
        self.scene.add_box(table_id, table_pose, table_size)

        # 设置目标物体的长宽高 [l, w, h]
        target_size = [0.05, 0.05, 0.23]
        # 设置目标的位姿，让目标物体位于桌子上
        target_pose = PoseStamped()
        target_pose.header.frame_id = REFERENCE_FRAME
        target_pose.pose.position.x = 0.47
        target_pose.pose.position.y = 0.0
        target_pose.pose.position.z =  target_size[2] / 2.0
        target_pose.pose.orientation.w = 1.0
        # 把目标物体添加到规划场景中
        self.scene.add_box(target_id, target_pose, target_size)

        rospy.sleep(1)

        # 设置桌子table为抓取和放置操作的支撑面，使MoveIt!忽略物体放到桌子上时产生的碰撞警告
        self.xarm.set_support_surface_name(table_id)

        # 初始化抓取的目标位姿
        grasp_pose = target_pose
        grasp_pose.pose.position.z += 0.02
        q = quaternion_from_euler(0, 0, 0)
        grasp_pose.pose.orientation.x = q[0]
        grasp_pose.pose.orientation.y = q[1]
        grasp_pose.pose.orientation.z = q[2]
        grasp_pose.pose.orientation.w = q[3]

        # 设置靠近抓取点的方向的向量[x,y,z]
        pre_grasp_approach_direction = [1.0, 0 , 0]
        # 生成grasp抓取列表
        grasps = self.make_grasps(grasp_pose, [table_id], pre_grasp_approach_direction)

        # 发布grasp_pose，用以在RViz中查看
        for grasp in grasps:
            self.gripper_pose_pub.publish(grasp.grasp_pose)
            rospy.sleep(0.2)

        # 设置result标记每次抓取尝试的结果
        result = None
        n_attempts = 0

        # 循环尝试抓取操作直到抓取成功或最大尝试次数用尽
        while result != MoveItErrorCodes.SUCCESS and n_attempts < max_pick_attempts:
            n_attempts += 1
            rospy.loginfo("Pick attempt: " +  str(n_attempts))
            result = self.xarm.pick(target_id, grasps)
            rospy.sleep(0.2)


        # 设置一个放置目标位姿
        place_pose = PoseStamped()
        place_pose.header.frame_id = REFERENCE_FRAME
        place_pose.pose.position.x = 0.32
        place_pose.pose.position.y = -0.32
        place_pose.pose.position.z = target_size[2] / 2.0
        q = quaternion_from_euler(0, 0, -math.pi/4)
        place_pose.pose.orientation.x = q[0]
        place_pose.pose.orientation.y = q[1]
        place_pose.pose.orientation.z = q[2]
        place_pose.pose.orientation.w = q[3]

        # 如果抓取成功，尝试物品放置操作。
        if result == MoveItErrorCodes.SUCCESS:
            # 重置标记操作结果的result和尝试次数
            result = None
            n_attempts = 0

            # 生成一系列放置位姿
            places = self.make_places(place_pose)

            # 循环放置直到放置成功或超过最大尝试次数
            while result != MoveItErrorCodes.SUCCESS and n_attempts < max_place_attempts:
                n_attempts += 1
                rospy.loginfo("Place attempt: " +  str(n_attempts))
                for place in places:
                    result = self.xarm.place(target_id, place)
                    if result == MoveItErrorCodes.SUCCESS:
                        break
                rospy.sleep(0.2)

            if result != MoveItErrorCodes.SUCCESS:
                rospy.loginfo("Place operation failed after " + str(n_attempts) + " attempts.")
                return CallPickPlaceDemoResponse(False)
        else:
            rospy.loginfo("Pick operation failed after " + str(n_attempts) + " attempts.")
            return CallPickPlaceDemoResponse(False)


        # 设置新的抓取位姿为放置点
        grasp_pose = place_pose
        grasp_pose.pose.position.z += 0.02
        q = quaternion_from_euler(0, 0, -math.pi/4)
        grasp_pose.pose.orientation.x = q[0]
        grasp_pose.pose.orientation.y = q[1]
        grasp_pose.pose.orientation.z = q[2]
        grasp_pose.pose.orientation.w = q[3]
        # 生成抓取位姿列表
        pre_grasp_approach_direction = [0, 0 , -1.0]
        grasps = self.make_grasps(grasp_pose, [table_id], pre_grasp_approach_direction)
        # 发布grasp_pose，用以在RViz中查看
        for grasp in grasps:
            self.gripper_pose_pub.publish(grasp.grasp_pose)
            rospy.sleep(0.2)

        # 设置result标记每次抓取尝试的结果
        result = None
        n_attempts = 0

        # 尝试抓取
        while result != MoveItErrorCodes.SUCCESS and n_attempts < max_pick_attempts:
            n_attempts += 1
            rospy.loginfo("Pick attempt: " +  str(n_attempts))
            result = self.xarm.pick(target_id, grasps)
            rospy.sleep(0.2)

        # 设置新的放置点为一开始物体放置的位置
        place_pose.pose.position.x = 0.47
        place_pose.pose.position.y = 0
        place_pose.pose.position.z = target_size[2] / 2.0
        q = quaternion_from_euler(0, 0, 0)
        place_pose.pose.orientation.x = q[0]
        place_pose.pose.orientation.y = q[1]
        place_pose.pose.orientation.z = q[2]
        place_pose.pose.orientation.w = q[3]
        # 若抓取成功，开始放置操作
        if result == MoveItErrorCodes.SUCCESS:
            result = None
            n_attempts = 0

            # Generate valid place poses
            places = self.make_places(place_pose)

            # Repeat until we succeed or run out of attempts
            while result != MoveItErrorCodes.SUCCESS and n_attempts < max_place_attempts:
                n_attempts += 1
                rospy.loginfo("Place attempt: " +  str(n_attempts))
                for place in places:
                    result = self.xarm.place(target_id, place)
                    if result == MoveItErrorCodes.SUCCESS:
                        break
                rospy.sleep(0.2)

            if result != MoveItErrorCodes.SUCCESS:
                rospy.loginfo("Place operation failed after " + str(n_attempts) + " attempts.")
                return CallPickPlaceDemoResponse(False)
        else:
            rospy.loginfo("Pick operation failed after " + str(n_attempts) + " attempts.")
            return CallPickPlaceDemoResponse(False)


        # 回到初始位姿
        self.xarm.set_named_target('Home')
        self.xarm.go()
        # 闭合手爪
        self.gripper.set_joint_value_target(GRIPPER_CLOSED)
        self.gripper.go()

        rospy.sleep(1)
        # Remove leftover objects from a previous run
        self.scene.remove_world_object(table_id)
        self.scene.remove_world_object(target_id)
        self.scene.remove_attached_object(GRIPPER_FRAME, target_id)
        rospy.sleep(1)
        # 返回服务的响应success为True
        return CallPickPlaceDemoResponse(True)


    # 通过手爪关节的位置生成JointTrajectory消息类型的轨迹返回值
    def make_gripper_posture(self, joint_positions):
        # Initialize the joint trajectory for the gripper joints
        t = JointTrajectory()

        # Set the joint names to the gripper joint names
        t.joint_names = GRIPPER_JOINT_NAMES

        # Initialize a joint trajectory point to represent the goal
        tp = JointTrajectoryPoint()

        # Assign the trajectory joint positions to the input positions
        tp.positions = joint_positions

        # Set the gripper effort
        tp.effort = GRIPPER_EFFORT

        tp.time_from_start = rospy.Duration(5)

        # Append the goal point to the trajectory points
        t.points.append(tp)

        # Return the joint trajectory
        return t

    # 通过传入的向量vector和最小距离、期望距离，生成GripperTranslation消息类型并返回
    def make_gripper_translation(self, min_dist, desired, vector):
        # Initialize the gripper translation object
        g = GripperTranslation()

        # Set the direction vector components to the input
        g.direction.vector.x = vector[0]
        g.direction.vector.y = vector[1]
        g.direction.vector.z = vector[2]

        # The vector is relative to the gripper frame
        g.direction.header.frame_id = REFERENCE_FRAME

        # Assign the min and desired distances from the input
        g.min_distance = min_dist
        g.desired_distance = desired
        return g

    # 生成一系列可能的抓取列表
    def make_grasps(self, initial_pose_stamped, allowed_touch_objects, pre_grasp_approach_direction):
        # 初始化抓取列表
        grasps = []
        # 把变量g初始化为Grasp消息类型的对象
        g = Grasp()

        # 设置 pre_grasp_posture和grasp_posture
        g.pre_grasp_posture = self.make_gripper_posture(GRIPPER_OPEN)
        g.grasp_posture = self.make_gripper_posture(GRIPPER_GRASP)

        # 设置靠近和撤离的方向以及距离
        g.pre_grasp_approach = self.make_gripper_translation(0.1, 0.12, pre_grasp_approach_direction)
        g.post_grasp_retreat = self.make_gripper_translation(0.08, 0.1, [0.0, 0.0, 1.0])

        # 设置grasp_pose
        g.grasp_pose = initial_pose_stamped

        # 设置这一个grasp的ID
        g.id = str(len(grasps))

        # 设置抓取过程中允许触碰的物体
        g.allowed_touch_objects = allowed_touch_objects

        # max_contact_force = 0 表示不设置抓握时要使用的最大接触力
        g.max_contact_force = 0

        # 把此grasp添加到grasps列表
        grasps.append(deepcopy(g))

        # 返回
        return grasps

    # 生成一系列可能的放置点位姿
    def make_places(self, init_pose):
        # Initialize the place location as a PlaceLocation message
        place = PlaceLocation()
        # Start with the input place pose
        place.place_pose = init_pose
        # 设置靠近放置点的方向、最小移动距离和期望距离，这里设置为沿着Z轴向下移动0.1米
        place.pre_place_approach.direction.header.frame_id = REFERENCE_FRAME;
        place.pre_place_approach.direction.vector.z = -1.0;
        place.pre_place_approach.min_distance = 0.08;
        place.pre_place_approach.desired_distance = 0.1;
        # 设置放置完成后机械臂的撤离方向、移动最小距离和期望距离，这里设置为沿着Z轴向上移动0.15米
        place.post_place_retreat.direction.header.frame_id = REFERENCE_FRAME;
        place.post_place_retreat.direction.vector.z = 1.0;
        place.post_place_retreat.min_distance = 0.12;
        place.post_place_retreat.desired_distance = 0.15;

        # A list of x shifts (meters) to try
        x_vals = [0, 0.005, 0.01, -0.005, -0.01]

        # A list of y shifts (meters) to try
        y_vals = [0, 0.005, 0.01, -0.005, -0.01]

        # A list to hold the places
        places = []

        # Generate a place pose for each angle and translation
        for y in y_vals:
            for x in x_vals:
                place.place_pose.pose.position.x = init_pose.pose.position.x + x
                place.place_pose.pose.position.y = init_pose.pose.position.y + y

                # Append this place pose to the list
                places.append(deepcopy(place))

        # Return the list
        return places

if __name__ == "__main__":
    MoveItPickPlaceDemo()
